import librosa
import librosa.display
import matplotlib.pyplot as plt
import numpy as np
import os

# 定义音频文件列表
audio_files = ['1-1.wav', '1-2.wav']
output_dir = 'output_1'  # 输出图片的文件夹

# 创建输出文件夹（如果不存在）
os.makedirs(output_dir, exist_ok=True)

# 处理每个音频文件
for audio_file in audio_files:
    # 加载音频文件
    y, sr = librosa.load(audio_file, sr=None)

    # 获取文件名
    file_name = os.path.splitext(os.path.basename(audio_file))[0]

    # 绘制波形图并保存
    plt.figure(figsize=(10, 4))
    librosa.display.waveshow(y, sr=sr)
    plt.title(f'Waveform - {file_name}')
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude')
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, f'{file_name}_waveform.png'))
    plt.close()

    # 绘制声谱图并保存
    plt.figure(figsize=(10, 4))
    S = librosa.stft(y)  # 计算短时傅里叶变换
    D = librosa.amplitude_to_db(np.abs(S), ref=np.max)  # 取绝对值后转换为分贝
    librosa.display.specshow(D, sr=sr, x_axis='time', y_axis='log')
    plt.colorbar(format='%+2.0f dB')
    plt.title(f'Spectrogram - {file_name}')
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, f'{file_name}_spectrogram.png'))
    plt.close()

    # 绘制 MFCC 特征并保存
    plt.figure(figsize=(10, 4))
    mfccs = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
    librosa.display.specshow(mfccs, sr=sr, x_axis='time')
    plt.colorbar()
    plt.title(f'MFCC - {file_name}')
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, f'{file_name}_mfcc.png'))
    plt.close()

print(f"处理完成，图片已保存到文件夹: {output_dir}")
